﻿import numpy as np
import math
TINY = 1.0e-27
def matrix(nrow,ncol):
    a=np.arange(nrow*ncol)
    return a.reshape(nrow,ncol)
def vector(nrow):
    return np.arange(nrow)
def lfit(x, y, p,  ndat,  ma):#, void (*funcs)(double, double [], int))
    a=np.arange(ma)
    chisq=0.0
    i=0
    j=0
    k=0
    l=0
    m=0
    ym=0.0
    wt=0.0
    sum1=0.0
    if ma == 0:
        return
    beta = matrix(ma, 1)
    afunc = vector(ma)
    covar = matrix(ma, ma)

    # for (j = 0; j < ma; j++):
    #     for (k = 0; k < ma; k++):
    #         covar[j][k] = 0.0;
    #     beta[j][0] = .0;
    for i in range(ndat):# (i = 0; i < ndat; i++):
        dpoly(x[i], afunc, ma);
        ym = y[i];
        j=-1
        l=0
        while(l<ma):
        # for (j = -1, l = 0; l < ma; l++):
            wt = afunc[l] * p[i];
            k=-1
            m=0
            j+=1
            while(m<=l):
            # for (j++, k = -1, m = 0; m <= l; m++):
                k+=1
                covar[j][k] += wt * afunc[m];
                m+=1
            beta[j][0] += ym * wt;
            l+=1
    for j in range(ma):#(j = 1; j < ma; j++):
        for k in range(j):#for (k = 0; k <= j; k++):
            covar[k][j] = covar[j][k]
    gaussj(covar, ma, beta, 1);
    j=-1
    l=0
    while(l<ma):
    # for (j = -1, l = 0; l < ma; l++):
        j+=1
        a[l] = beta[j][0]
        l+=1
    chisq = 0.0;
    for i in range(ndat):#(i = 0; i < ndat; i++):
        dpoly(x[i], afunc, ma)
        sum1=0.0
        j=0
        while(j<ma):
        # for (sum1 = .0, j = 0; j < ma; j++):
            sum1 += a[j] * afunc[j];
            j+=1
        chisq += math.pow(y[i] - sum1, 2.0);

    return (a,chisq)
def gaussj(a,n,b,m):#double[][] a, int n, double[][] b, int m)
    indxc=0
    indxr=0
    ipiv=0
    i=0
    icol = 0
    irow = 0
    j=0
    k=0
    l=0
    ll=0
    big=0.0
    dum=0.0
    pivinv=0.0

    indxc = np.arange(n)
    indxr = np.arange(n)
    ipiv = np.arange(n)

    for j in range(n):#(j = 0; j < n; j++)
        ipiv[j] = -1;

    for i in range(n):#(i = 0; i < n; i++)
        big = .0;
        for j in range(n):#(j = 0; j < n; j++)
            if (ipiv[j] != 0):
                for k in range(n):# (k = 0; k < n; k++)
                    if (ipiv[k] == -1):
                        if (math.fabs(a[j][k]) >= big):
                            big = math.fabs(a[j][k]);
                            irow = j;
                            icol = k;
                    else:
                        if (ipiv[k] > 0):
                            return;
        ipiv[icol]+=1
        if (irow != icol):
            for l in range(n):#(l = 0; l < n; l++)
                tmp=a[irow][l]
                a[irow][l]=a[icol][l]
                a[icol][l]=tmp
            for l in range(n):#(l = 0; l < m; l++)
                tmp=b[irow][l]
                b[irow][l]=b[icol][l]
                b[icol][l]=tmp
        indxr[i] = irow;
        indxc[i] = icol;
        if (a[icol][icol] == .0):
            return;
        pivinv = 1.0 / a[icol][icol];
        a[icol][icol] = 1.0;
        for l in range(n):# (l = 0; l < n; l++)
            a[icol][l] *= pivinv;
        for l in range(m):# (l = 0; l < m; l++)
            b[icol][l] *= pivinv;
        for ll in range(n):#(ll = 0; ll < n; ll++)
            if (ll != icol):
                dum = a[ll][icol];
                a[ll][icol] = .0;
                for l in range(n):# (l = 0; l < n; l++)
                    a[ll][l] -= a[icol][l] * dum;
                for l in range(m):# (l = 0; l < m; l++)
                    b[ll][l] -= b[icol][l] * dum;

    l=n-1
    while(l>=0):
    # for (l = n - 1; l >= 0; l--)
        if (indxr[l] != indxc[l]):
            for k in range(n):#for (k = 0; k < n; k++)
                tmp=a[k][indxr[l]]
                a[k][indxr[l]]=a[k][indxc[l]]
                a[k][indxc[l]]=tmp
        l-=1

def dpoly(x,p,np):#double x, double[] p, int np):
    p[0] = 1.0;
    for j in range(np):#(int j = 1; j < np; j++):
        p[j] = p[j - 1] * x;
if __name__=="__main__":
    x=[1,2,3,4]
    y=[2,4,6,8]
    p=[1,1,1,1]
    r=lfit(x,y,p,len(x),2)
    print(r)
